
#include "../include/SSS_SETUP.h"
//#include <__clang_cuda_device_functions.h>
#include <cuda_runtime.h>
#include <cuda_runtime_api.h>
#include <driver_types.h>
#include <curand.h>
#include <helper_cuda.h>

#include <nsparse.cuh>
#include <nsparse_asm.h>

double SoC_time = 0;
double Aggregates_time = 0;
double Tentative_time = 0;
double Interp_time = 0;
double Trans_time = 0;
double RAP_time = 0;

// double SpGEMM_count = 0;
// double SpMV_count =0 ;

double metis_time = 0;



void SSS_Setup_print(double Setup_time,double SoC_time,double Aggregates_time,double Tentative_time,double Interp_time,double Trans_time,double RAP_time)
{
    printf("\n");
    printf("Setup_time       = %.2lf ms          \n",Setup_time);
    printf("---------------------Real Time: ---------------------\n");
    printf("SoC_time         = %.2lf ms          \n",SoC_time);
    printf("Aggregates_time  = %.2lf ms          \n",Aggregates_time);
    printf("Tentative_time   = %.2lf ms          \n",Tentative_time);
    printf("Interp_time      = %.2lf ms          \n",Interp_time);
    printf("Trans_time       = %.2lf ms          \n",Trans_time);
    printf("RAP_time         = %.2lf ms          \n",RAP_time);

    double SoC = SoC_time/Setup_time *100;
    double Aggregates= Aggregates_time/Setup_time *100;
    double Tentative= Tentative_time/Setup_time *100;
    double Interp = Interp_time/Setup_time *100;
    double Trans = Trans_time/Setup_time *100;
    double RAP = RAP_time/Setup_time *100;

    
    printf("\n");
    printf("---------------------Component/Setup ratio: ---------------------\n");
    printf("SoC             /  Setup   = %.2lf %%      \n",SoC);
    printf("Aggregates      /  Setup   = %.2lf %%      \n",Aggregates);
    printf("Tentative       /  Setup   = %.2lf %%      \n",Tentative);
    printf("Interp          /  Setup   = %.2lf %%      \n",Interp);
    printf("Trans           /  Setup   = %.2lf %%      \n",Trans);
    printf("RAP             /  Setup   = %.2lf %%      \n",RAP);

    printf("\n");
}


void SSS_amg_complexity_print(SSS_AMG *mg)
{
    const int max_levels = mg->num_levels;
    int level;
    double gridcom = 0.0, opcom = 0.0;

    printf("-----------------------------------------------------------\n");
    printf("  Level   Num of rows   Num of nonzeros   Avg. NNZ / row   \n");
    printf("-----------------------------------------------------------\n");

    for (level = 0; level < max_levels; ++level) {
        double AvgNNZ = (double) mg->cg[level].A.num_nnzs / mg->cg[level].A.num_rows;

        printf("%5d %13d %17d %14.2lf\n", level, mg->cg[level].A.num_rows,
                mg->cg[level].A.num_nnzs, AvgNNZ);

        gridcom += mg->cg[level].A.num_rows;
        opcom += mg->cg[level].A.num_nnzs;
    }

    printf("-----------------------------------------------------------\n");

    gridcom /= mg->cg[0].A.num_rows;
    opcom /= mg->cg[0].A.num_nnzs;
    printf("  Grid complexity = %.3lf  |", gridcom);
    printf("  Operator complexity = %.3lf\n", opcom);

    printf("-----------------------------------------------------------\n");
}




sfCSR spgemm_csr(sfCSR *a, sfCSR *b, sfCSR *c)
{
    /* Execution of SpGEMM on Device */


    spgemm_kernel_hash(a, b, c);

    cudaDeviceSynchronize();
    return *c;

}


void SSS2sfCSR(const SSS_MAT *src, sfCSR *des)
{
    des->M = src->num_rows;
    des->N = src->num_cols;
    des->nnz = src->num_nnzs;

    des->d_rpt = src->d_row_ptr;
    des->d_col = src->d_col_idx;
    des->d_val = src->d_val;
}

void sfCSR2SSS(const sfCSR *src, SSS_MAT *des)
{
    des->num_rows = src->M;
    des->num_cols = src->N;
    des->num_nnzs = src->nnz;

    des->d_row_ptr = src->d_rpt;
    des->d_col_idx = src->d_col;
    des->d_val = src->d_val;
}


SSS_MAT SSS_mat_rap_device( SSS_MAT *R,  SSS_MAT *A,  SSS_MAT *P)
{
    SSS_MAT tmp;
    SSS_MAT RAP;

    //cuSPARSE

    SSS_blas_spgemm_device(R,A,&tmp);
    SSS_blas_spgemm_device(&tmp,P,&RAP);

    //nsparse
    // printf("bengin nsparse\n");
    // printf("create sfCSR\n");
    // sfCSR mat_R, mat_A, mat_P;
    // sfCSR mat_tmp,mat_RAP;
    
    // printf("sss2sfCSR\n");
    // SSS2sfCSR(R,&mat_R);
    // SSS2sfCSR(A,&mat_A);
    // SSS2sfCSR(P,&mat_P);

    // printf("nsparse spgemm\n");
    // spgemm_csr(&mat_R, &mat_A, &mat_tmp);
    // printf("nsparse RAP end \n");
    // spgemm_csr(&mat_tmp, &mat_P, &mat_RAP);
    
    // printf("sfCSR2sss\n");
    // sfCSR2SSS(&mat_RAP,&RAP);
    
    // printf("nsparse end\n");

    return RAP;

}



void SSS_amg_setup(SSS_AMG *mg, SSS_MAT *A, SSS_AMG_PARS *pars)
{ 
    int min_cdof; //最小粗变量数

    int m;
    int status = 0;
    int lvl = 0, max_lvls;
    double setup_start, setup_end;

    SSS_IMAT S_MAT;   //强连接矩阵
    int size;

    min_cdof = SSS_max(pars->coarse_dof, MIN_CDOF);  //最小粗变量数
    printf("min_cdof = %d\n",min_cdof);
    max_lvls = pars->max_levels;

    /* timer */
    setup_start = SSS_get_time();

    cudaError_t EE;
    /* create mg */
    //create mg struct
    *mg = SSS_amg_data_create(pars);
    

    int num_candidates = 1;
    
    SSS_VEC B = SSS_vec_create(A->num_rows);
    
    B = SSS_vec_create_device(A->num_rows,B);
    
    for(int i=0;i< A->num_rows;i++)
    {
        B.d[i]=1.0;
    }

    EE = cudaMemcpy(B.d_d ,B.d  ,A->num_rows * sizeof(double),cudaMemcpyHostToDevice);
    if(EE != cudaSuccess)
    {
        printf("B to device faild\n");
    }
    
    /* init */
    m = A->num_rows;

    /* init matrix */
    //Create CSR sparse matrix data memory space
    
    //给 cg[0].A 开device空间
    mg->cg[0].A = SSS_mat_struct_create_device(m, m, A->num_nnzs,mg->cg[0].A);

    EE = cudaMemcpy(mg->cg[0].A.d_row_ptr,A->row_ptr, (A->num_rows+1) * sizeof(int), cudaMemcpyHostToDevice);
    if(EE != cudaSuccess)
    {
        printf("mg->cg[0].A row_ptr fail \n");
    }
    EE = cudaMemcpy(mg->cg[0].A.d_col_idx,A->col_idx, A->num_nnzs * sizeof(int), cudaMemcpyHostToDevice);
    if(EE != cudaSuccess)
    {
        printf("mg->cg[0].A col_idx fail \n");
    }
    EE = cudaMemcpy(mg->cg[0].A.d_val,A->val, A->num_nnzs * sizeof(double), cudaMemcpyHostToDevice);
    if(EE != cudaSuccess)
    {
        printf("mg->cg[0].A val fail \n");
    }

    printf(" ---------------- Main AMG setup loop ----------------\n");
    

    // Main AMG setup loop
    // 当前层粗网格的行数 > 最小粗节点数  或  当前层 < 最大层数 -1
    printf("max_lvls = %d\n",max_lvls);
    while ((mg->cg[lvl].A.num_rows > min_cdof) && (lvl < max_lvls - 1)) 
    {

        printf("=====================================\n");
        printf("==================lvl = %d=================\n",lvl);
        printf("mg->cg[%d].A: rows = %d, cols = %d , nnz = %d\n",lvl,mg->cg[lvl].A.num_rows,mg->cg[lvl].A.num_cols,mg->cg[lvl].A.num_nnzs);

        mg->cg[lvl].A.d_dig = SSS_mat_get_diag_GPU(&mg->cg[lvl].A,0);

        //SoC
        printf("(1.)SOC\n");
        double SoC1 = SSS_get_time();
        status = SoC_GPU(&mg->cg[lvl].A, &S_MAT, pars);
        double SoC2 = SSS_get_time();
        SoC_time +=SoC2 - SoC1;
        printf("SSS_calloc???\n");

        //临时开空间,晓涵metis需要
        S_MAT.row_ptr = (int *) malloc((S_MAT.num_rows+1) * sizeof(int));
        S_MAT.col_idx = (int *) malloc(S_MAT.num_nnzs * sizeof(int));
        
        printf("SSS_calloc!!!\n");
        EE = cudaMemcpy(S_MAT.row_ptr,S_MAT.d_row_ptr, (S_MAT.num_rows+1) * sizeof(int), cudaMemcpyDeviceToHost);
        if(EE != cudaSuccess)
        {
            printf("S_MAT row_ptr fail \n");
        }
        EE = cudaMemcpy(S_MAT.col_idx,S_MAT.d_col_idx, S_MAT.num_nnzs * sizeof(int), cudaMemcpyDeviceToHost);
        if(EE != cudaSuccess)
        {
            printf("S_MAT col_idx fail \n");
        }

        
        //聚集
        printf("(2.)聚集\n");
        double aggregates1 = SSS_get_time();
        SSS_IVEC aggregates = SSS_amg_agg_GPU(&mg->cg[lvl].A, &mg->cg[lvl].P, &S_MAT,pars);     //status为0代表粗化成功
        double aggregates2 = SSS_get_time();
        Aggregates_time +=aggregates2 - aggregates1;


        //临时插值T
        printf("(3.)临时插值T\n");
        double Tentative1 = SSS_get_time();
        SSS_MAT T = tentative_p_GPU(&aggregates,&B,pars);    
        double Tentative2 = SSS_get_time();
        Tentative_time +=Tentative2 - Tentative1;


        //Jacobi插值P
        printf("(4.)Jacobi插值P\n");
        double  Interp1 = SSS_get_time();
        SSS_amg_jacobi_prolongation_GPU(&mg->cg[lvl].A, &T, &mg->cg[lvl].P); 
        double  Interp2 = SSS_get_time();
        Interp_time +=Interp2 - Interp1;

      


        //Check : Is coarse sparse too small?
        if (mg->cg[lvl].P.num_cols < min_cdof) 
        {
            printf("check 2\n");
            /*-- Clean up S_MAT generated in coarsening --*/
            SSS_free(S_MAT.row_ptr);
            SSS_free(S_MAT.col_idx);

            break;
        }

        // Check : Does this coarsening step too aggressive?
        if (mg->cg[lvl].P.num_rows > mg->cg[lvl].P.num_cols * 10) 
        {
                printf("### WARNING: Coarsening might be too aggressive!\n");
                printf("### WARNING: Lvl = %d ,Fine level = %d, coarse level = %d. Discard!\n",
                     lvl,mg->cg[lvl].P.num_rows, mg->cg[lvl].P.num_cols);
        }
       
        printf("(5.)trans R\n");
        double Trans1 = SSS_get_time();
        mg->cg[lvl].R = SSS_mat_trans_csr2csc_device(&mg->cg[lvl].P);
        double Trans2 = SSS_get_time();
        Trans_time +=Trans2 - Trans1;

        printf("(6.)rap\n");
        double RAP1 = SSS_get_time();

        //mg->cg[lvl + 1].A = SSS_blas_mat_rap_device(&mg->cg[lvl].R, &mg->cg[lvl].A, &mg->cg[lvl].P);
        mg->cg[lvl + 1].A = SSS_mat_rap_device(&mg->cg[lvl].R, &mg->cg[lvl].A, &mg->cg[lvl].P);
        
        double RAP2 = SSS_get_time();
        RAP_time +=RAP2 - RAP1;
               
        /*-- Clean up S_MAT generated in coarsening --*/
        SSS_free(S_MAT.row_ptr);
        SSS_free(S_MAT.col_idx);

        cudaFree(S_MAT.d_row_ptr);
        cudaFree(S_MAT.d_col_idx);

        // Check : Is the coarse matrix too dense?
        if (mg->cg[lvl].A.num_nnzs / mg->cg[lvl].A.num_rows > mg->cg[lvl].A.num_cols * 0.2) 
        {
                printf("### WARNING: Coarse matrix is too dense!\n");
                printf("### WARNING: m = n = %d, nnz = %d!\n", mg->cg[lvl].A.num_cols,
                        mg->cg[lvl].A.num_nnzs);
            /* free A */
            //SSS_mat_destroy(&mg->cg[lvl + 1].A);
            break;
        }

        lvl++;
    }

    // setup total level number and current level
    mg->num_levels = max_lvls = lvl + 1;
    
    //mg->cg[0].wp = SSS_vec_create(m);
    //开device上的空间
    mg->cg[0].wp = SSS_vec_create_device(m,mg->cg[0].wp);

    //给每一层的x和b开device上的空间
    //mg->cg[0].b 和 x 在Solve开（最细层就是原始b x） 
    for (lvl = 1; lvl < max_lvls; ++lvl)
    {
        int mm = mg->cg[lvl].A.num_rows;
        mg->cg[lvl].b = SSS_vec_create_device(mm,mg->cg[lvl].b);
        mg->cg[lvl].x = SSS_vec_create_device(mm,mg->cg[lvl].x);
        mg->cg[lvl].wp = SSS_vec_create_device(2 * mm,mg->cg[lvl].wp);
    }

    setup_end = SSS_get_time();

    SSS_amg_complexity_print(mg);
    
    double Setup_time = setup_end - setup_start;

    printf("metis_time         = %.2lf ms          \n",metis_time);
    
    SSS_Setup_print(Setup_time,SoC_time,Aggregates_time,Tentative_time,Interp_time,Trans_time,RAP_time);


}